

import librosa
import os
import numpy as np
from scipy.io import wavfile

raw_path = './dataset/test_offline/task3_estimate_raw'
target_path = './dataset/test_offline/task3_estimate'
try:
    os.mkdir(target_path)
except:
    print('Already Exists')

raw_list = os.listdir(raw_path)
raw_list.sort()


for audio_idx in range(len(raw_list)):

    audio = raw_list[audio_idx]
    audio_path = os.path.join(raw_path, audio)
    print(audio)

    if audio_idx % 3 == 0:
        audio1_path = audio_path.replace('left', 'middle')
        audio2_path = audio_path.replace('left', 'right')
    elif audio_idx % 3 == 1:
        audio1_path = audio_path.replace('middle', 'left')
        audio2_path = audio_path.replace('middle', 'right')
    elif audio_idx % 3 == 2:
        audio1_path = audio_path.replace('right', 'left')
        audio2_path = audio_path.replace('right', 'middle')
    target_audio_path = os.path.join(target_path, audio)

    # norm_audio_path = audio_path.replace('.wav','-norm.wav')
    # norm_audio1_path = audio1_path.replace('.wav','-norm.wav') 
    # norm_audio2_path = audio2_path.replace('.wav','-norm.wav') 

    # os.system('ffmpeg-normalize '+audio_path+' -ar 44100 -o '+norm_audio_path+' -f')
    # os.system('ffmpeg-normalize '+audio1_path+' -ar 44100 -o '+norm_audio1_path+' -f')
    # os.system('ffmpeg-normalize '+audio2_path+' -ar 44100 -o '+norm_audio2_path+' -f')

    y, sr = librosa.load(audio_path, sr=44100)
    y1, sr1 = librosa.load(audio1_path, sr=44100)
    y2, sr2 = librosa.load(audio2_path, sr=44100)

    # norm_factor0 = np.max(np.abs(y)) * 1.1
    # norm_factor1 = np.max(np.abs(y1)) * 1.1
    # norm_factor2 = np.max(np.abs(y2)) * 1.1

    # norm_factor = max([norm_factor0, norm_factor1, norm_factor2])

    # y /= norm_factor0
    # y1 /= norm_factor1
    # y2 /= norm_factor2
    
    y_out = y - y1*0.2 - y2*0.2

    librosa.output.write_wav(target_audio_path, y_out, 44100)

os.system('rm -rf ./dataset/test_offline/task3_estimate_raw/*-norm.wav')